package com.ai.model.xunfei;

import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import okhttp3.Response;
import okhttp3.WebSocket;
import okhttp3.WebSocketListener;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;
import org.springframework.util.CollectionUtils;
import org.springframework.util.StringUtils;

import java.util.UUID;

/**
 * @author hh
 * @since 2024/5/26
 */
public class XfChatAiWebsocket extends WebSocketListener {

    private final XfApiProperties xfApiProperties;
    private final String newQuestion;
    private final StringBuilder answer = new StringBuilder();
    private int status;
    private static final int FINISHED = 2;

    public XfChatAiWebsocket(XfApiProperties xfApiProperties, String newQuestion) {
        this.xfApiProperties = xfApiProperties;
        this.newQuestion = newQuestion;
    }

    @Override
    public void onClosed(@NotNull WebSocket webSocket, int code, @NotNull String reason) {
        super.onClosed(webSocket, code, reason);
    }

    @Override
    public void onClosing(@NotNull WebSocket webSocket, int code, @NotNull String reason) {
        super.onClosing(webSocket, code, reason);
    }

    @Override
    public void onFailure(@NotNull WebSocket webSocket, @NotNull Throwable t, @Nullable Response response) {
        super.onFailure(webSocket, t, response);
    }

    @Override
    public void onMessage(@NotNull WebSocket webSocket, @NotNull String responseText) {
//        System.out.println(responseText);
        JSONObject resultJson = JSONObject.parseObject(responseText);
        JSONObject payload = resultJson.getJSONObject("payload");
        JSONObject choices = payload.getJSONObject("choices");
        synchronized (this) {
            status = choices.getIntValue("status");
            JSONArray text = choices.getJSONArray("text");
            if (CollectionUtils.isEmpty(text)) {
                return;
            }
            for (Object o : text) {
                JSONObject textObj = (JSONObject) o;
                String role = textObj.getString("role");
                String content = textObj.getString("content");
                if ("assistant".equals(role) && StringUtils.hasText(content)) {
                    answer.append(content);
                }
            }
        }
    }

    @Override
    public void onOpen(@NotNull WebSocket webSocket, @NotNull Response response) {
        String requestParameter = getRequestParameter(newQuestion);
        webSocket.send(requestParameter);
    }

    private String getRequestParameter(String newQuestion) {
        JSONObject requestJson = new JSONObject();
        // header参数
        JSONObject headerJson = new JSONObject();
        headerJson.put("app_id", xfApiProperties.getAppId());
        headerJson.put("uid", UUID.randomUUID().toString().substring(0, 10));
        // parameter参数
        JSONObject parameterJson = new JSONObject();
        JSONObject chatJson = new JSONObject();
        chatJson.put("domain", xfApiProperties.getDomain());
        chatJson.put("temperature", 0.5);
        chatJson.put("max_tokens", 4096);
        parameterJson.put("chat", chatJson);
        // payload参数
        JSONObject payloadJson = new JSONObject();
        JSONObject messageJson = new JSONObject();
        JSONArray textArray = new JSONArray();
        JSONObject oneTextJson = new JSONObject();
        oneTextJson.put("role", "user");
        oneTextJson.put("content", newQuestion);
        textArray.add(oneTextJson);
        messageJson.put("text", textArray);
        payloadJson.put("message", messageJson);

        requestJson.put("header", headerJson);
        requestJson.put("parameter", parameterJson);
        requestJson.put("payload", payloadJson);
        return requestJson.toJSONString();
    }

    public String getAnswer() {
        return answer.toString();
    }

    public boolean isFinished() {
        return status == FINISHED;
    }

}
